"""
机载相机数据接口，处理机载相机视频流，截取符合要求的帧

date: 2025/9/21
author: SiHeng Tang
file: camera.py
copyright(c) DFSA Eastern Software
此程序不能提供任何担保 WITHOUT WARRANTY OF ANY KIND
"""
import base64
import logging
import queue
import time
from copy import deepcopy
from multiprocessing import Process, Queue

import cv2
import numpy as np

from async_sche import TaskWrapper, BaseTaskGen
from mengji import Vehicle
from mod_loader import ConfigHandler
from task_element import DataEndpoint, DataPack

logger = logging.getLogger(__name__)
handler = ConfigHandler(__name__)

DEFAULT_CONFIG = {
    "label": "camera",
    "device": "/dev/video0",
    "img_width": 1280,
    "img_height": 800,
    "fps": 30,
    "drop_frame_frac": 4,
    "max_pipe_cache": 3,
    "encoding_type": ".webp",
    "encoding_quality": 2
}


def frame2string(_f: np.ndarray, _q: int = 10, _t: str = ".webp") -> str:
    """
    将 opencv 图像转换为 base64 字符串，用于 RPC 传输
    :param _f: opencv 图像
    :param _q: 图像质量 0-100 默认 100
    :param _t: 压缩算法，默认 webp
    :return: base64 字符串
    """
    if _t == ".jpg":
        cv_option = cv2.IMWRITE_JPEG_QUALITY
    else:
        cv_option = cv2.IMWRITE_WEBP_QUALITY

    ok, buffer = cv2.imencode(_t, _f, [cv_option, _q])

    if not ok:
        raise RuntimeError("image encode failed")

    return base64.b64encode(buffer.tobytes()).decode('ascii')


def string2frame(_s: str) -> np.ndarray:
    """
    将 base64 编码的 JPEG 图像转换为 opencv 图像
    :param _s: base64 字符串
    :return: opencv 图像，即 ndarray
    """
    raw_array = np.frombuffer(base64.b64decode(_s), np.uint8)
    cv_image = cv2.imdecode(raw_array, cv2.IMREAD_COLOR)

    if cv_image is None:
        raise RuntimeError("image decode failed")

    return cv_image


class ProcessedFetch(Process):
    """
    进程化图像采集，避免阻塞主线程运行
    """

    def __init__(self, config: dict, _q: Queue):
        super().__init__(None, None, None, daemon=True)

        self._dev_path: str = config["device"]
        self._img_size: tuple[int, int] = (config["img_width"], config["img_height"])
        self._fps: int = config["fps"]
        self._drop_frame_frac: int = config["drop_frame_frac"]
        self._encoding: str = config["encoding_type"]
        self._encoding_qlt: int = config["encoding_quality"]

        self._is_open = False
        self._capture: cv2.VideoCapture | None = None
        self._capture_cnt: int = 0

        self._frame_out = _q

    def cat_one(self):
        """从视频流中抓取一帧图像"""
        if not self._is_open:
            logger.warning(f"try open cam: {self._dev_path}")
            self._capture = cv2.VideoCapture(self._dev_path)

            if self._capture.isOpened():
                self._capture.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'))
                self._capture.set(cv2.CAP_PROP_FRAME_WIDTH, self._img_size[0])
                self._capture.set(cv2.CAP_PROP_FRAME_HEIGHT, self._img_size[1])
                self._capture.set(cv2.CAP_PROP_FPS, self._fps)
                self._is_open = True

        ret, frame = self._capture.read()
        timestamp = time.time()

        if not ret:
            logger.error(f"capture camera failed, is camera disconnected?")
            self._capture.release()
            self._is_open = False

            # 避免过于平凡重试
            time.sleep(0.5)
            return

        self._capture_cnt += 1

        if not (self._capture_cnt % self._drop_frame_frac):
            return

        timed_frame = {
            "frame": frame,
            "cap_time": timestamp,
            "str_frame": frame2string(frame, self._encoding_qlt, self._encoding)
        }

        try:
            self._frame_out.put(deepcopy(timed_frame), block=False)
        except queue.Full:
            self._frame_out.get()
            self._frame_out.put(deepcopy(timed_frame))

        self._capture_cnt = 0

    def run(self):
        while True:
            self.cat_one()


class OpencvCamera(BaseTaskGen):
    """
    通过 opencv 接口采集视频流转换为帧，加载时间标签并发送给其他处理节点
    相机采集速度完全是根据帧率和丢帧比例确定的
    """

    def __init__(self, config: dict):
        super().__init__(priority=TaskWrapper.IDLE_PRIORITY + 10)

        self._frame_in = Queue(config["max_pipe_cache"])
        self._cat = ProcessedFetch(config, self._frame_in)
        self._cat.start()

        self.pub_timed_frame = DataEndpoint()

    async def _func(self):
        try:
            timed_frame = self._frame_in.get(block=False)
        except queue.Empty:
            return
        else:
            self.pub_timed_frame.set_dp(DataPack(payload=timed_frame))

    def _on_exit(self):
        self.pub_timed_frame = None


@handler.add_handler("virtual_bus:camera")
def load_config(veh: Vehicle, cfg_section: dict):
    logger.debug(f"optics scanner load config {cfg_section}")

    cfg = deepcopy(DEFAULT_CONFIG)
    cfg.update(cfg_section)

    veh.reg_part(cfg["label"], OpencvCamera(cfg), is_sche=True)
